#include <fstream>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <TChain.h>
#include <TFile.h>
#include <TMath.h>

void ge_dc_precond(std::string const &);
void turn(double, double, double, double, double *, double *);

void ge_dc_precond(std::string const &file_name)
{
  TChain *in = new TChain("is430");
  in->Add(file_name.c_str());
  if (0 == in->GetEntries()) {
    std::cerr << "Problem opening input file.\n";
    return;
  }

  std::ifstream angle_file("ToolsCalibration/cluster_angles.txt");
  if (!angle_file.is_open()) {
    throw std::runtime_error("Could not open angles file.");
  }
  double start_angles[8][4];
  for (int i = 0; 8 > i; ++i) {
    char line[80];
    angle_file.getline(line, 80);
    std::istringstream iss(line);
    for (int j = 0; 4 > j; ++j) {
      float f;
      iss >> f;
      start_angles[i][j] = f;
    }
  }
  angle_file.close();
  double th[8][24], ph[8][24];
  for (int clu = 0; 8 > clu; ++clu) {
    turn(start_angles[clu][0], start_angles[clu][1], start_angles[clu][2],
        start_angles[clu][3], th[clu], ph[clu]);
  }

  float in_ge_e;
  float in_ge_dc_e;
  Long64_t in_ge_t;
  UShort_t in_ge_id;
  in->SetBranchAddress("ge_e", &in_ge_e);
  in->SetBranchAddress("ge_dc_e", &in_ge_dc_e);
  in->SetBranchAddress("ge_t", &in_ge_t);
  in->SetBranchAddress("ge_id", &in_ge_id);

  float in_fb_de;
  float in_fb_e;
  Long64_t in_fb_t;
  UChar_t in_fb_det_id;
  UChar_t in_fb_strip_id;
  float in_fb_strip_pos;
  float in_fb_theta;
  float in_fb_phi;
  float in_fb_p_de;
  float in_fb_p_e;
  in->SetBranchAddress("fb_de", &in_fb_de);
  in->SetBranchAddress("fb_e", &in_fb_e);
  in->SetBranchAddress("fb_t", &in_fb_t);
  in->SetBranchAddress("fb_det_id", &in_fb_det_id);
  in->SetBranchAddress("fb_strip_id", &in_fb_strip_id);
  in->SetBranchAddress("fb_strip_pos", &in_fb_strip_pos);
  in->SetBranchAddress("fb_theta", &in_fb_theta);
  in->SetBranchAddress("fb_phi", &in_fb_phi);
  in->SetBranchAddress("fb_p_de", &in_fb_p_de);
  in->SetBranchAddress("fb_p_e", &in_fb_p_e);

  size_t dot_idx = file_name.find_last_of('.');
  if (0 != file_name.substr(dot_idx).compare(".root")) {
    std::cerr << "Filename ending is weard.\n";
    return;
  }
  std::string out_file_name = file_name.substr(0, dot_idx) +
      "_dc_precond.root";
  TFile file_out(out_file_name.c_str(), "recreate");
  if (!file_out.IsOpen()) {
    std::cerr << "Problem opening output file.\n";
    return;
  }
  TTree *out = new TTree("is430", "is430");

  float ge_e;
  UShort_t ge_id;
  float ge_theta;
  float ge_phi;
  out->Branch("ge_e", &ge_e, "ge_e/F");
  out->Branch("ge_id", &ge_id, "ge_id/s");
  out->Branch("ge_theta", &ge_theta, "ge_theta/F");
  out->Branch("ge_phi", &ge_phi, "ge_phi/F");

  UChar_t fb_channel;
  float fb_beta;
  float fb_gamma;
  float fb_x;
  float fb_y;
  float fb_z;
  out->Branch("fb_channel", &fb_channel, "fb_channel/b");
  out->Branch("fb_beta", &fb_beta, "fb_beta/F");
  out->Branch("fb_gamma", &fb_gamma, "fb_gamma/F");
  out->Branch("fb_x", &fb_x, "fb_x/F");
  out->Branch("fb_y", &fb_y, "fb_y/F");
  out->Branch("fb_z", &fb_z, "fb_z/F");

  int map[] = {1, 2, 4, 5, 6, 7};
  int entry_num = in->GetEntries();
  std::cout << entry_num << '\n';
  for (int i = 0; entry_num > i; ++i) {
    in->GetEntry(i);
    if (100>abs(in_fb_t-in_ge_t-290)) {
      continue;
    }
    if (400>in_fb_e) {
      continue;
    }
    // dE-E cuts:
    //  Ne:
    //   Upper.: 3.8e7, 4.5e3
    //   Middle: 2.2e7, 3.7e3
    //   Lower.: 8.4e6, 2e3
    //  Be:
    //   Triton..: 4.2e7, 3.4e3
    //   Deuteron: 3.2e7, 3.8e3
    //   Proton..: 1.9e7, 3e3
    //   Bottom..: 9.8e6, 2.2e3
/*    if (in_fb_de < 8.4e6 / (in_fb_e + 2.0e3)) {
      continue;
    }
    else if (in_fb_de < 2.2e7 / (in_fb_e + 3.7e3)) {
      fb_channel = 1;
    }
    else if (in_fb_de < 3.8e7 / (in_fb_e + 4.5e3)) {
      fb_channel = 2;
    }*/
    if (in_fb_de < 9.8e6 / (in_fb_e + 2.2e3)) {
      continue;
    }
    else if (in_fb_de < 1.9e7 / (in_fb_e + 3e3)) {
      fb_channel = 1;
    }
    else if (in_fb_de < 3.2e7 / (in_fb_e + 3.8e3)) {
      fb_channel = 2;
    }
    else if (in_fb_de < 4.2e7 / (in_fb_e + 3.4e3)) {
      fb_channel = 3;
    }
    else {
      continue;
    }
    // 950 - 1016 - 1070.
/*    if (900>in_ge_e || 1140<in_ge_e) {
      continue;
    }*/

    // Transform proton angles to neon angles.
    float const kPi = 3.14159265;
    float t = in_fb_theta / kPi * 180.0;
    float ne_theta = (
        1.2575e-01 + (
        5.6046e-04 + (
        -4.7355e-05 + (
        5.0427e-07 + (
        -2.1773e-09 + 
        3.4529e-12 * t
        ) * t
        ) * t
        ) * t
        ) * t
        ) * t * kPi / 180.0;
    float ne_phi = fmod(in_fb_phi + kPi, 2 * kPi);
    float ne_e =
        1.7510e+00 + (
        -4.7549e-04 + (
        1.7075e-04 + (
        -1.5010e-06 +
        3.7628e-09 * t
        ) * t
        ) * t
        ) * t;

    ge_e = in_ge_e;
    ge_id = in_ge_id;
    int clu = ge_id / 100;
    int cry = (ge_id - 100 * clu) / 10;
    int seg = ge_id - 10 * cry - 100 * clu;
    ge_theta = th[clu][cry * 8 + map[seg]];
    ge_phi = ph[clu][cry * 8 + map[seg]];
    fb_gamma = ne_e / 939.0 + 1;
    fb_beta = sqrt(1 - 1 / (fb_gamma * fb_gamma));
    fb_x = -sin(ne_phi) * sin(ne_theta);
    fb_y = cos(ne_phi) * sin(ne_theta);
    fb_z = cos(ne_theta);

    out->Fill();
  }
  file_out.Write();
  file_out.Close();
}

void turn(double t, double p, double a, double dist, double *th, double *ph) 
{
  double const kDeg2Rad = TMath::Pi() / 180.0;
  t = t * kDeg2Rad;
  p = p * kDeg2Rad;
  a = a * kDeg2Rad;
  double const kCore2Crystal = 34.0;

  double center_x[3];
  double center_y[3];
  double center_z[3];

  // Core A.
  center_x[0] = kCore2Crystal;
  center_y[0] = 0;
  center_z[0] = dist;

  // Core B.
  center_x[1] = -kCore2Crystal * sin(30 * kDeg2Rad);
  center_y[1] = -kCore2Crystal * cos(30 * kDeg2Rad);
  center_z[1] = dist;

  // Core C.
  center_x[2] = -kCore2Crystal * sin(30 * kDeg2Rad);
  center_y[2] = kCore2Crystal * cos(30 * kDeg2Rad);
  center_z[2] = dist;

  double const kR = kCore2Crystal * 0.6;
  double const kRs = kR * sin(30.0 * kDeg2Rad);
  double const kRc = kR * cos(30.0 * kDeg2Rad);

  double x[24];
  double y[24];
  double z[24];

  // Crystal A.
  x[0] = center_x[0];        // A0 (core)
  y[0] = center_y[0];
  z[0] = center_z[0];
  x[1] = center_x[0] + kRs;  // A1
  y[1] = center_y[0] - kRc;
  z[1] = center_z[0];
  x[2] = center_x[0] - kRs;  // A2
  y[2] = center_y[0] - kRc;
  z[2] = center_z[0];
  x[3] = 0.0;                // empty DGF channel
  y[3] = 0.0;
  z[3] = 0.0;
  x[4] = center_x[0] + kR;   // A3
  y[4] = center_y[0] + 0.0;
  z[4] = center_z[0];
  x[5] = center_x[0] + kRs;  // A4
  y[5] = center_y[0] + kRc;
  z[5] = center_z[0];
  x[6] = center_x[0] - kR;   // A5
  y[6] = center_y[0] + 0.0;
  z[6] = center_z[0];
  x[7] = center_x[0] - kRs;  // A6
  y[7] = center_y[0] + kRc;
  z[7] = center_z[0];

  // Crystal B.
  x[8] = center_x[1];
  y[8] = center_y[1];
  z[8] = center_z[1];
  x[9] = center_x[1] - kR;
  y[9] = center_y[1] + 0.0;
  z[9] = center_z[1];
  x[10] = center_x[1] - kRs;
  y[10] = center_y[1] + kRc;
  z[10] = center_z[1];
  x[11] = 0.0;
  y[11] = 0.0;
  z[11] = 0.0;
  x[12] = center_x[1] - kRs;
  y[12] = center_y[1] - kRc;
  z[12] = center_z[1];
  x[13] = center_x[1] + kRs;
  y[13] = center_y[1] - kRc;
  z[13] = center_z[1];
  x[14] = center_x[1] + kRs;
  y[14] = center_y[1] + kRc;
  z[14] = center_z[1];
  x[15] = center_x[1] + kR;
  y[15] = center_y[1] + 0.0;
  z[15] = center_z[1];

  // Crystal C.
  x[16] = center_x[2];
  y[16] = center_y[2];
  z[16] = center_z[2];
  x[17] = center_x[2] + kRs;
  y[17] = center_y[2] + kRc;
  z[17] = center_z[2];
  x[18] = center_x[2] + kR;
  y[18] = center_y[2] + 0.0;
  z[18] = center_z[2];
  x[19] = 0.0;
  y[19] = 0.0;
  z[19] = 0.0;
  x[20] = center_x[2] - kRs;
  y[20] = center_y[2] + kRc;
  z[20] = center_z[2];
  x[21] = center_x[2] - kR;
  y[21] = center_y[2] + 0.0;
  z[21] = center_z[2];
  x[22] = center_x[2] + kRs;
  y[22] = center_y[2] - kRc;
  z[22] = center_z[2];
  x[23] = center_x[2] - kRs;
  y[23] = center_y[2] - kRc;
  z[23] = center_z[2];

  // matrix = m_phi * m_theta * m_alpha
  double const cp = cos(p);
  double const sp = sin(p);
  double const ct = cos(t);
  double const st = sin(t);
  double const ca = cos(a);
  double const sa = sin(a);
  // Column major.
  double matrix[9] = {
    ca * cp - ct * sa * sp, cp * ct * sa + ca * sp, -sa * st,
    -cp * sa - ca * ct * sp, ca * cp * ct - sa * sp, -ca * st,
    -sp * st, cp * st, ct
  };

  // Transform!
  for (int i = 0; 24 > i; ++i) {
    double xi = matrix[0] * x[i] + matrix[3] * y[i] + matrix[6] * z[i];
    double yi = matrix[1] * x[i] + matrix[4] * y[i] + matrix[7] * z[i];
    double zi = matrix[2] * x[i] + matrix[5] * y[i] + matrix[8] * z[i];
    th[i] = atan2(sqrt(xi * xi + yi * yi), zi);
    ph[i] = atan2(-xi, yi);
  }
}
